from typing import Callable, List, Optional, Union
from MelodieFuncFlow import MelodieGenerator, MelodieFrozenGenerator
from .. import universal_ast_nodes as nodes
from ..unparser import BaseUASTUnparser


# class FiltersCreator:
def create_method_name_filter(
    method_name: str,
) -> Callable[[nodes.SourceElement], bool]:
    def wrapper(method: nodes.SourceElement):
        match method:
            case nodes.MethodDecl(name=nodes.Name(id=name_id)):
                return name_id == method_name
            case _:
                return False

    return wrapper


class UASTQuery:
    @classmethod
    def get_all_globals(
        cls, uast: nodes.SourceElement, exclude_externs: bool = False
    ) -> List[nodes.VarDecl]:
        """
        Get all globals from one file
        """
        all_globals = []

        def search_global_def_in_uast(ctx: nodes.SourceElement.ApplyContext) -> bool:
            """
            Search global definitions in the UAST
            """
            if ctx.current_node.__class__ in {nodes.MethodDecl}:
                return False

            if isinstance(ctx.current_node, nodes.VarDecl):
                if exclude_externs and ctx.current_node.is_extern:
                    return False
                else:
                    all_globals.append(ctx.current_node)
            return True

        uast.apply_with_hierarchy(search_global_def_in_uast)
        return all_globals

    @classmethod
    def get_all_params(cls, uast: nodes.MethodDecl) -> list[nodes.ParamDecl]:
        """
        Get all parameter from the method or function
        """
        return uast.filter_by(nodes.ParamDecl).l

    @classmethod
    def get_all_locals(
        cls, uast: Union[nodes.MethodDecl, nodes.NameSpaceDef]
    ) -> List[nodes.VarDecl]:
        """
        Get all local variable declarations from uast
        uast should be:
        """
        return uast.filter_by(nodes.VarDecl).l

    @classmethod
    def get_all_params(
        cls, uast: nodes.MethodDecl
    ) -> MelodieGenerator[nodes.ParamDecl]:
        """
        Get all parameter from the method or function
        """
        return uast.filter_by(nodes.ParamDecl)

    @classmethod
    def get_method(
        cls,
        uast: nodes.SourceElement,
        method_name: Optional[str] = None,
        definition_only: bool = True,
    ) -> nodes.MethodDecl:
        """
        通过指定的条件，获取函数或方法的声明节点，并返回第一个

        :method_name: 函数或方法的名称
        :return: MethodDecl
        """
        try:
            return (
                cls.iter_methods(uast, method_name)
                .filter(
                    lambda decl: (decl.body is not None) if definition_only else True
                )
                .head()
            )
        except StopIteration:
            raise ValueError(f"No method named `{method_name}` found in this uast")

    @classmethod
    def iter_methods(
        cls,
        uast: nodes.SourceElement,
        method_name: Optional[str] = None,
        definition_only: bool = True,
    ) -> MelodieGenerator[nodes.MethodDecl]:
        """
        通过指定的条件，获取函数或方法的声明节点们，并返回一个生成器

        :method_name: 函数或方法的名称
        :return: MethodDecl
        """

        methods_generator = uast.filter_by(nodes.MethodDecl)
        if method_name is not None:
            methods_generator = methods_generator.filter(
                create_method_name_filter(method_name)
            )
        if definition_only:
            methods_generator = methods_generator.filter(
                lambda decl: decl.body is not None
            )
        return methods_generator

    @classmethod
    def iter_global_assignments(
        cls, node: nodes.SourceElement
    ) -> MelodieFrozenGenerator[str]:
        """
        Iterate all assigned global variable names in the given node.

        """
        all_assigned_variables = cls.iter_variable_assignments(node).s
        all_locals = {local.variable.id for local in cls.get_all_locals(node)}
        all_params = {param.name.id for param in cls.get_all_params(node)}
        return MelodieFrozenGenerator(
            list(all_assigned_variables - (all_locals | all_params))
        )

    @classmethod
    def iter_variable_assignments(
        cls,
        node: nodes.SourceElement,
    ) -> MelodieGenerator[str]:
        """
        Get all assigned variable names in the given node.
        Note that "VarDecl" nodes with initializer will not be included.

        :param node: The node to search for variable assignments.
        """

        def get_assigned_names(node: nodes.SourceElement):
            match node:
                case nodes.Name(id=variable_name):
                    return variable_name
                case nodes.ArrayAccessExpr(target=target):
                    return get_assigned_names(target)
                case nodes.FieldAccessExpr(target=target):
                    return get_assigned_names(target)
                case nodes.ReferenceExpr(value=value):
                    return get_assigned_names(value)
                case nodes.UnaryExpr(sign=sign, expression=expression):
                    return get_assigned_names(expression)
                case nodes.CastExpr(expression=expression):
                    return get_assigned_names(expression)
                case nodes.TypeRef() | nodes.CallExpr() | nodes.BinaryExpr():
                    return BaseUASTUnparser().unparse(node)
                case _:
                    raise ValueError(f"Unsupported assignment type: {node}")

        def _():
            for assigned_item in node.filter_by(nodes.Assignment).map(
                lambda assignment: assignment.lhs
            ):
                assert len(assigned_item) == 1
                yield get_assigned_names(assigned_item[0])

        return MelodieGenerator(_())
